import torch
import torchvision
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import os

data_dir = './data'
model_dir = ['./checkpoint/B0',
             './checkpoint/B1',
             './checkpoint/B2',
             './checkpoint/B3']

# 图像预处理
mean = [0.5070751592371323, 0.48654887331495095, 0.4409178433670343]
std = [0.2673342858792401, 0.2564384629170883, 0.27615047132568404]
batch_size = 128  # 原文是1024
num_workers = 4  # 2

# 原文提到了random erase的处理方法 但是有网页指出 这个方法对CIFAR-100没什么帮助
# 参考网页：https://blog.csdn.net/u013685264/article/details/122564323
transform_train = transforms.Compose([
    # transforms.Resize((224, 224)),
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(15),
    transforms.ToTensor(),
    transforms.Normalize(mean, std),
    transforms.RandomErasing()
])
transform_test = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize(mean, std)])

# CIFAR-100 数据集下载

data_name = './data/cifar-100-python'
# if not os.path.isdir(data_name):
#     train_dataset = torchvision.datasets.CIFAR100(root='data/',
#                                                   train=True,
#                                                   transform=transform_train,
#                                                   download=True)
# else:
#     train_dataset = torchvision.datasets.CIFAR100(root='data/',
#                                                   train=True,
#                                                   transform=transform_train,
#                                                   download=False)
train_dataset = torchvision.datasets.CIFAR100(root='data/',
                                              train=True,
                                              transform=transform_train,
                                              download=False)

test_dataset = torchvision.datasets.CIFAR100(root='data/',
                                             train=False,
                                             transform=transform_test)

# 数据载入

train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                           batch_size=batch_size,
                                           num_workers=num_workers,
                                           shuffle=True)

test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                          batch_size=batch_size,
                                          num_workers=num_workers,
                                          shuffle=False)
